#!/bin/bash
# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

############################################################################
## This script generates the model repository needed by some of the
## tritonserver CI tests. Generating these models requires using
## the PyTorch container.
##
## 1. Update TENSORRT_IMAGE and PYTORCH_IMAGE to
## match what is being used by the tritonserver release being
## tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## local system that you want to target for the generated models.
##
## 3. Run this script to create /tmp/qa_model_repository,
## /tmp/qa_variable_model_repository, /tmp/qa_shapetensor_model_repository
## /tmp/qa_identity_model_repository, /tmp/qa_identity_big_model_repository
## /tmp/qa_reshape_model_repository, /tmp/qa_noshape_model_repository,
## /tmp/qa_sequence_model_repository, /tmp/qa_ensemble_model_repository
## /tmp/qa_dyna_sequence_model_repository, and
## /tmp/qa_variable_sequence_model_repository directories containing
## all the models needed for CI testing.
##
############################################################################
TRITON_MDLS_BASE_SCRIPT_DIR="$(dirname $(readlink -f $0))"
TRITON_MDLS_BASE_SCRIPT_FILE="$(readlink -f $0)"

log_message.info() { local message=$@ ; echo -e "\033[34m$(date +"%Y-%m-%d %H:%M:%S") - [ INFO ] - ${message} \033[0m"; } ;
log_message.status() { local message=$@ ; echo -e "\033[37m$(date +"%Y-%m-%d %H:%M:%S") - [ STATUS ] - ${message} \033[0m"; } ;
log_message.warning() { local message=$@ ; echo -e "\033[33m$(date +"%Y-%m-%d %H:%M:%S") - [ WARNING ] - ${message} \033[0m"; } ;
log_message.error() { local message=$@ ; echo -e "\033[31m$(date +"%Y-%m-%d %H:%M:%S") - [ ERROR ] - ${message} \033[0m"; } ;

log_message.status "Changing working directory to the script directory to: " "${TRITON_MDLS_BASE_SCRIPT_DIR}"

cd ${TRITON_MDLS_BASE_SCRIPT_DIR}

log_message.status "define: default values"
TRITON_VERSION=${TRITON_VERSION:=25.11}
ONNX_VERSION=1.16.1
ONNX_OPSET=0
OPENVINO_VERSION=2024.5.0
UBUNTU_IMAGE=${UBUNTU_IMAGE:=ubuntu:22.04}
PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
TENSORRT_IMAGE=${TENSORRT_IMAGE:=nvcr.io/nvidia/tensorrt:$TRITON_VERSION-py3}
CUDA_DEVICE=${NV_GPU:=0}
MODEL_TYPE=${MODEL_TYPE:-""}
TRITON_MODELS_USE_DOCKER=${TRITON_MODELS_USE_DOCKER:-"1"}
TRITON_MODELS_USE_ENROOT=${TRITON_MODELS_USE_ENROOT:-"1"}

log_message.status "define: CI related parameters"
CI_JOB_ID=${CI_JOB_ID:=$(date +%Y%m%d_%H%M)}
RUNNER_ID=${RUNNER_ID:=0}
PROJECT_NAME=${PROJECT_NAME:=tritonserver}

log_message.status "define: Docker engine parameters"
DOCKER_VOLUME=${DOCKER_VOLUME:=volume.gen_qa_model_repository.${CI_JOB_ID}}
DOCKER_VOLUME_CONTAINER=${DOCKER_VOLUME}.gen_qa_model_repository.${CI_JOB_ID}
DOCKER_GPU_ARGS=${DOCKER_GPU_ARGS:-$([[ $RUNNER_GPUS =~ ^[0-9] ]] && eval $NV_DOCKER_ARGS || echo "--gpus device=$CUDA_DEVICE" )}

function define_model_output_directories() {
    log_message.status "define: model output directories"
    TRITON_MDLS_BLD_DIR=${TRITON_MDLS_BLD_DIR:=/mnt/$CI_JOB_ID}
    TRITON_MDLS_SRC_DIR=${TRITON_MDLS_SRC_DIR:=$TRITON_MDLS_BLD_DIR/gen_srcdir}
    TRITON_MDLS_QA_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_model_repository
    TRITON_MDLS_QA_VARIABLE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_variable_model_repository
    TRITON_MDLS_QA_IDENTITY_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_identity_model_repository
    TRITON_MDLS_QA_IDENTITY_BIG_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_identity_big_model_repository
    TRITON_MDLS_QA_SHAPETENSOR_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_shapetensor_model_repository
    TRITON_MDLS_QA_RESHAPE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_reshape_model_repository
    TRITON_MDLS_QA_SEQUENCE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_sequence_model_repository
    TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_dyna_sequence_model_repository
    TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_dyna_sequence_implicit_model_repository
    TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_variable_sequence_model_repository
    TRITON_MDLS_QA_ENSEMBLE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_ensemble_model_repository
    TRITON_MDLS_QA_NOSHAPE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_noshape_model_repository
    TRITON_MDLS_QA_TRT_PLUGIN_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_trt_plugin_model_repository
    TRITON_MDLS_QA_RAGGED_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_ragged_model_repository
    TRITON_MDLS_QA_TRT_FORMAT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_trt_format_model_repository
    TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_trt_data_dependent_model_repository
    TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_sequence_implicit_model_repository
    TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_variable_sequence_implicit_model_repository
    TRITON_MDLS_QA_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_sequence_initial_state_implicit_model_repository
    TRITON_MDLS_QA_VARIABLE_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_variable_sequence_initial_state_implicit_model_repository
    TRITON_MDLS_QA_TORCHTRT_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/torchtrt_model_store
    TRITON_MDLS_QA_SCALAR_MODELS=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_scalar_models
    TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_dynamic_batch_image_model_repository
    TRITON_MDLS_QA_CUSTOM_OPS=$TRITON_MDLS_BLD_DIR/$TRITON_VERSION/qa_custom_ops/libtorch_custom_ops
}

function define_models_generation_scripts() {
    log_message.status "define: script names"
    SCRIPT_NAME_SUFFIX=${SCRIPT_NAME_SUFFIX:-v2}

    OPENVINOSCRIPT=gen.OpenVINO.gen_qa_model_repository.${SCRIPT_NAME_SUFFIX}.sh
    log_message.info "OpenVINO script: " "${OPENVINOSCRIPT}"

    ONNXSCRIPT=gen.ONNXRuntime.gen_qa_model_repository.${SCRIPT_NAME_SUFFIX}.sh
    log_message.info "ONNX script: " "${ONNXSCRIPT}"

    TORCHSCRIPT=gen.PyTorch.gen_qa_model_repository.${SCRIPT_NAME_SUFFIX}.sh
    log_message.info "PyTorch script: " "${TORCHSCRIPT}"

    TRTSCRIPT=gen.TensorRT.gen_qa_model_repository.${SCRIPT_NAME_SUFFIX}.sh
    log_message.info "TensorRT script: " "${TRTSCRIPT}"

    log_message.status "create: OpenVINO script - ${OPENVINOSCRIPT}"
    cat > $OPENVINOSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential \
        cmake \
        libprotobuf-dev \
        protobuf-compiler \
        python3 \
        python3-dev \
        python3-pip \
        wget \
        gnupg2 \
        software-properties-common

ln -s /usr/bin/python3 /usr/bin/python

pip3 install  "numpy<=1.23.5" setuptools

pip3 install openvino==$OPENVINO_VERSION

# Since variable shape tensors are not allowed, identity models may fail to generate.
# TODO Add variable size tensor models after DLIS-2827 adds support for variable shape tensors.
# TODO Add sequence models after DLIS-2864 adds support for sequence/control inputs.
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --openvino --models_dir=$TRITON_MDLS_QA_MODEL
chmod -R 777 $TRITON_MDLS_QA_MODEL
# python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --openvino --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
# chmod -R 777 $TRITON_MDLS_QA_IDENTITY_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_reshape_models.py --openvino --models_dir=$TRITON_MDLS_QA_RESHAPE_MODEL
chmod -R 777 $TRITON_MDLS_QA_RESHAPE_MODEL
# python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --openvino --models_dir=$TRITON_MDLS_QA_SEQUENCE_MODEL
# chmod -R 777 $SVOLUME_EQDESTDIR
# python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --openvino --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
# chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
EOF

    log_message.status "run: chmod a+x ${OPENVINOSCRIPT}"
    chmod a+x $OPENVINOSCRIPT

    if [ $? -ne 0 ]; then
        log_message.error "failed: chmod ${OPENVINOSCRIPT}"
        exit 1
    fi

    log_message.status "create: ONNX script - ${ONNXSCRIPT}"
    cat > $ONNXSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
        apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
                protobuf-compiler python3 python3-dev python3-pip
ln -s /usr/bin/python3 /usr/bin/python

pip3 install "protobuf<=3.20.1"  "numpy<=1.23.5" # TODO: Remove current line DLIS-3838
pip3 install --upgrade onnx==${ONNX_VERSION}

python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_MODEL
chmod -R 777 $TRITON_MDLS_QA_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
chmod -R 777 $TRITON_MDLS_QA_IDENTITY_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_reshape_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$TRITON_MDLS_QA_RESHAPE_MODEL
chmod -R 777 $TRITON_MDLS_QA_RESHAPE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_ragged_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_RAGGED_MODEL
chmod -R 777 $TRITON_MDLS_QA_RAGGED_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_ort_scalar_models.py --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_SCALAR_MODELS
chmod -R 777 $TRITON_MDLS_QA_SCALAR_MODELS
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --ensemble --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --ensemble --variable --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_variable_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_reshape_models.py --ensemble --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_reshape_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --ensemble --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_identity_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --ensemble --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_sequence_model_repository
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --ensemble --variable --models_dir=$TRITON_MDLS_QA_ENSEMBLE_MODEL/qa_variable_sequence_model_repository
chmod -R 777 $TRITON_MDLS_QA_ENSEMBLE_MODEL
EOF

    log_message.status "run: chmod a+x ${ONNXSCRIPT}"
    chmod a+x ${ONNXSCRIPT}
    if [ $? -ne 0 ]; then
        log_message.error "failed: chmod ${ONNXSCRIPT}"
        exit 1
    fi

    log_message.status "create: PyTorch script - ${TORCHSCRIPT}"
    cat > $TORCHSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
pip3 install onnxscript
set -e
set -x
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --models_dir=$TRITON_MDLS_QA_MODEL
chmod -R 777 $TRITON_MDLS_QA_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --libtorch --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
chmod -R 777 $TRITON_MDLS_QA_IDENTITY_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_reshape_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_RESHAPE_MODEL
chmod -R 777 $TRITON_MDLS_QA_RESHAPE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --libtorch --models_dir=$TRITON_MDLS_QA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --libtorch --models_dir=$TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
if [ -z "$MODEL_TYPE" ] || [ "$MODEL_TYPE" != "igpu" ]; then
  nvidia-smi --query-gpu=compute_cap | grep -qz 12.1 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for Torch TensorRT' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_torchtrt_models.py --models_dir=$TRITON_MDLS_QA_TORCHTRT_MODEL
  chmod -R 777 $TRITON_MDLS_QA_TORCHTRT_MODEL
fi
python3 $TRITON_MDLS_SRC_DIR/gen_qa_ragged_models.py --libtorch --models_dir=$TRITON_MDLS_QA_RAGGED_MODEL
chmod -R 777 $TRITON_MDLS_QA_RAGGED_MODEL
# Export torchvision image models to ONNX
python3 $TRITON_MDLS_SRC_DIR/gen_qa_image_models.py --resnet50 --resnet152 --vgg19 --models_dir=$TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL

export TORCH_EXTENSIONS_DIR=/tmp/.cache/torch_extensions/
mkdir -p \${TORCH_EXTENSIONS_DIR}
python3 $TRITON_MDLS_SRC_DIR/gen_qa_custom_ops_models.py --libtorch --models_dir=$TRITON_MDLS_QA_CUSTOM_OPS
mkdir -p $TRITON_MDLS_QA_CUSTOM_OPS/libtorch_modulo/
cp \${TORCH_EXTENSIONS_DIR}/custom_modulo/custom_modulo.so $TRITON_MDLS_QA_CUSTOM_OPS/libtorch_modulo/.
chmod -R 777 $TRITON_MDLS_QA_CUSTOM_OPS
EOF

    log_message.status "run: chmod a+x ${TORCHSCRIPT}"
    chmod a+x ${TORCHSCRIPT}
    if [ $? -ne 0 ]; then
        log_message.error "failed: chmod ${TORCHSCRIPT}"
        exit 1
    fi

    log_message.status "create: TensorRT script - ${TRTSCRIPT}"
    cat > $TRTSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
dpkg -l | grep TensorRT
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
# Models using shape tensor i/o
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$TRITON_MDLS_QA_SHAPETENSOR_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$TRITON_MDLS_QA_SHAPETENSOR_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$TRITON_MDLS_QA_SHAPETENSOR_MODEL
chmod -R 777 $TRITON_MDLS_QA_SHAPETENSOR_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_MODEL
chmod -R 777 $TRITON_MDLS_QA_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --tensorrt --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --tensorrt-compat --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
chmod -R 777 $TRITON_MDLS_QA_IDENTITY_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$TRITON_MDLS_QA_IDENTITY_BIG_MODEL
chmod -R 777 $TRITON_MDLS_QA_IDENTITY_BIG_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$TRITON_MDLS_QA_RESHAPE_MODEL
chmod -R 777 $TRITON_MDLS_QA_RESHAPE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_implicit_models.py --tensorrt --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_dyna_sequence_implicit_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL
chmod -R 777 $TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_ragged_models.py --tensorrt --models_dir=$TRITON_MDLS_QA_RAGGED_MODEL
chmod -R 777 $TRITON_MDLS_QA_RAGGED_MODEL
python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_format_models.py --models_dir=$TRITON_MDLS_QA_TRT_FORMAT_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_FORMAT_MODEL
nvidia-smi --query-gpu=compute_cap | grep -qz 11.0 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for data dependent shape' || python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_data_dependent_shape.py --models_dir=$TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL
# Make shared library for custom Hardmax plugin.
if [ -d "/usr/src/tensorrt/samples/python/onnx_custom_plugin" ]; then
    cd /usr/src/tensorrt/samples/python/onnx_custom_plugin
else
    git clone -b release/\${TRT_VERSION%.*.*} --depth 1 https://github.com/NVIDIA/TensorRT.git /workspace/TensorRT
    cd /workspace/TensorRT/samples/python/onnx_custom_plugin
fi
rm -rf build && mkdir build && \
cd build && cmake .. && make -j && cp libcustomHardmaxPlugin.so $TRITON_MDLS_QA_TRT_PLUGIN_MODEL/.
LD_PRELOAD=$TRITON_MDLS_QA_TRT_PLUGIN_MODEL/libcustomHardmaxPlugin.so python3 $TRITON_MDLS_SRC_DIR/gen_qa_trt_plugin_models.py --models_dir=$TRITON_MDLS_QA_TRT_PLUGIN_MODEL
chmod -R 777 $TRITON_MDLS_QA_TRT_PLUGIN_MODEL
EOF

    log_message.status "run: chmod a+x ${TRTSCRIPT}"
    chmod a+x ${TRTSCRIPT}
    if [ $? -ne 0 ]; then
        log_message.error "failed: chmod ${TRTSCRIPT}"
        exit 1
    fi

}

log_message.status "check: engine installation"
if [ "$TRITON_MODELS_USE_DOCKER" -eq 1 ] && which docker ; then
    log_message.info "Docker is installed."

    define_model_output_directories

    SCRIPT_NAME_SUFFIX=docker.v2 define_models_generation_scripts

    if ! docker volume inspect $DOCKER_VOLUME > /dev/null 2>&1; then
        log_message.status "docker volume: $DOCKER_VOLUME does not exist. Creating..."
        docker volume create $DOCKER_VOLUME --label RUNNER_ID=$RUNNER_ID --label PROJECT_NAME=$PROJECT_NAME
        log_message.status "docker volume: $DOCKER_VOLUME created"
        docker volume inspect $DOCKER_VOLUME
    else
        log_message.status "docker volume: $DOCKER_VOLUME in use"
        docker volume inspect $DOCKER_VOLUME
    fi

    log_message.status "docker pull: $UBUNTU_IMAGE"
    docker pull $UBUNTU_IMAGE

    log_message.status "docker volume: create destination directory on volume"
    docker run \
        --rm \
        --label RUNNER_ID=$RUNNER_ID \
        --label PROJECT_NAME=$PROJECT_NAME \
        -v $DOCKER_VOLUME:/mnt \
        -w /mnt/$CI_JOB_ID \
        $UBUNTU_IMAGE \
        mkdir -p gen_srcdir ${TRITON_VERSION}

    docker run \
        --rm \
        --label RUNNER_ID=$RUNNER_ID \
        --label PROJECT_NAME=$PROJECT_NAME \
        -v $DOCKER_VOLUME:/mnt \
        -w /mnt/$CI_JOB_ID \
        $UBUNTU_IMAGE \
        mkdir -p \
        $TRITON_MDLS_BLD_DIR \
        $TRITON_MDLS_SRC_DIR \
        $TRITON_MDLS_QA_MODEL \
        $TRITON_MDLS_QA_VARIABLE_MODEL \
        $TRITON_MDLS_QA_IDENTITY_MODEL \
        $TRITON_MDLS_QA_IDENTITY_BIG_MODEL \
        $TRITON_MDLS_QA_SHAPETENSOR_MODEL \
        $TRITON_MDLS_QA_RESHAPE_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_ENSEMBLE_MODEL \
        $TRITON_MDLS_QA_NOSHAPE_MODEL \
        $TRITON_MDLS_QA_TRT_PLUGIN_MODEL \
        $TRITON_MDLS_QA_RAGGED_MODEL \
        $TRITON_MDLS_QA_TRT_FORMAT_MODEL \
        $TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_TORCHTRT_MODEL \
        $TRITON_MDLS_QA_SCALAR_MODELS \
        $TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL

    log_message.status "docker container: create container $DOCKER_VOLUME_CONTAINER"
    docker create \
        --label RUNNER_ID=$RUNNER_ID \
        --label PROJECT_NAME=$PROJECT_NAME \
        --name $DOCKER_VOLUME_CONTAINER \
        -v $DOCKER_VOLUME:/mnt \
        -w /mnt/$CI_JOB_ID \
        $UBUNTU_IMAGE

    log_message.status "docker container: copy script to container"
    docker cp . $DOCKER_VOLUME_CONTAINER:/mnt/$CI_JOB_ID/gen_srcdir

    if [[ "aarch64" != $(uname -m) ]] ; then

        log_message.status "docker run: $OPENVINOSCRIPT"
        docker run \
            --rm \
            --label RUNNER_ID=$RUNNER_ID \
            --label PROJECT_NAME=$PROJECT_NAME \
            $DOCKER_GPU_ARGS \
            -v $DOCKER_VOLUME:/mnt \
            $UBUNTU_IMAGE \
            bash -xe $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT

        if [ $? -ne 0 ]; then
            log_message.error "docker run: ${OPENVINOSCRIPT} failed"
            exit 1
        fi
    fi # [[ "aarch64" != $(uname -m) ]]

    log_message.status "docker run: $ONNXSCRIPT"
    docker run \
        --rm \
        --label RUNNER_ID=$RUNNER_ID \
        --label PROJECT_NAME=$PROJECT_NAME \
        $DOCKER_GPU_ARGS \
        -v $DOCKER_VOLUME:/mnt \
        $UBUNTU_IMAGE \
        bash -xe $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT

    if [ $? -ne 0 ]; then
        log_message.error "docker run: ${ONNXSCRIPT} failed"
        exit 1
    fi

    log_message.status "docker pull: $PYTORCH_IMAGE"
    docker pull $PYTORCH_IMAGE

    log_message.status "docker run: $TORCHSCRIPT"
    docker run \
        --rm \
        --label RUNNER_ID=$RUNNER_ID \
        --label PROJECT_NAME=$PROJECT_NAME \
        $DOCKER_GPU_ARGS \
        -v $DOCKER_VOLUME:/mnt \
        $PYTORCH_IMAGE \
        bash -xe $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT

    if [ $? -ne 0 ]; then
        log_message.error "docker run: ${TORCHSCRIPT} failed"
        exit 1
    fi

    if [ "$MODEL_TYPE" != "igpu" ] ; then
        log_message.status "docker pull: $TENSORRT_IMAGE"
        docker pull $TENSORRT_IMAGE

        log_message.status "docker run: $TRTSCRIPT"
        docker run \
            --rm \
            --label RUNNER_ID=$RUNNER_ID \
            --label PROJECT_NAME=$PROJECT_NAME \
            $DOCKER_GPU_ARGS \
            -v $DOCKER_VOLUME:/mnt \
            -e TRT_VERBOSE \
            $TENSORRT_IMAGE \
            bash -xe $TRITON_MDLS_SRC_DIR/$TRTSCRIPT

        if [ $? -ne 0 ]; then
            log_message.error "docker run: ${TRTSCRIPT} failed"
            exit 1
        fi
    fi # [ "$MODEL_TYPE" != "igpu" ]

    if [ -z $CI ] ; then
        log_message.status "docker cp:copying generated models to /tmp/"
        docker cp $DOCKER_VOLUME_CONTAINER:$TRITON_MDLS_BLD_DIR/$TRITON_VERSION /tmp/
        log_message.status "docker rm: removing docker container $DOCKER_VOLUME_CONTAINER"
        docker rm -f $(docker ps -a --filter volume=$DOCKER_VOLUME --format '{{ .ID }}')
        log_message.status "docker volume rm: removing docker volume $DOCKER_VOLUME"
        docker volume rm $DOCKER_VOLUME
    fi # [ -z $CI ]

elif [ "$TRITON_MODELS_USE_ENROOT" -eq 1 ] && which enroot ; then
    log_message.info "NVIDIA Enroot is installed." ;

    TRITON_MDLS_BLD_DIR="/tmp/$CI_JOB_ID" define_model_output_directories

    SCRIPT_NAME_SUFFIX=enroot.v1 define_models_generation_scripts

    log_message.status "cleanup models folder if exists: $TRITON_MDLS_BLD_DIR"
    rm -rf $TRITON_MDLS_BLD_DIR

    log_message.status "create models directory structure in: $TRITON_MDLS_BLD_DIR"
     mkdir -p \
        $TRITON_MDLS_BLD_DIR \
        $TRITON_MDLS_SRC_DIR \
        $TRITON_MDLS_QA_MODEL \
        $TRITON_MDLS_QA_VARIABLE_MODEL \
        $TRITON_MDLS_QA_IDENTITY_MODEL \
        $TRITON_MDLS_QA_IDENTITY_BIG_MODEL \
        $TRITON_MDLS_QA_SHAPETENSOR_MODEL \
        $TRITON_MDLS_QA_RESHAPE_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_DYNA_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_DYNA_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_MODEL \
        $TRITON_MDLS_QA_ENSEMBLE_MODEL \
        $TRITON_MDLS_QA_NOSHAPE_MODEL \
        $TRITON_MDLS_QA_TRT_PLUGIN_MODEL \
        $TRITON_MDLS_QA_RAGGED_MODEL \
        $TRITON_MDLS_QA_TRT_FORMAT_MODEL \
        $TRITON_MDLS_QA_TRT_DATA_DEPENDENT_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_VARIABLE_SEQUENCE_INITIAL_STATE_IMPLICIT_MODEL \
        $TRITON_MDLS_QA_TORCHTRT_MODEL \
        $TRITON_MDLS_QA_SCALAR_MODELS \
        $TRITON_MDLS_QA_DYNAMIC_BATCH_IMAGE_MODEL

    log_message.status "copy scripts to: $TRITON_MDLS_SRC_DIR"
    cp -rv $TRITON_MDLS_BASE_SCRIPT_DIR/* $TRITON_MDLS_SRC_DIR/

    log_message.status "enroot import: $UBUNTU_IMAGE to ubuntu.$CI_JOB_ID.enroot.sqsh"
    enroot import --output /tmp/ubuntu.$CI_JOB_ID.enroot.sqsh docker://$UBUNTU_IMAGE


    log_message.status "enroot create: openvino.ubuntu.$CI_JOB_ID"
    enroot create --name openvino.ubuntu.$CI_JOB_ID /tmp/ubuntu.$CI_JOB_ID.enroot.sqsh
    log_message.info "enroot start: openvino.ubuntu.$CI_JOB_ID"
    enroot start --root --rw -m /tmp:/tmp openvino.ubuntu.$CI_JOB_ID bash -xe $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT


    log_message.status "enroot create: onnxruntime.ubuntu.$CI_JOB_ID"
    enroot create --name onnxruntime.ubuntu.$CI_JOB_ID /tmp/ubuntu.$CI_JOB_ID.enroot.sqsh
    log_message.info "enroot start: onnxruntime.ubuntu.$CI_JOB_ID"
    enroot start --root --rw -m /tmp:/tmp onnxruntime.ubuntu.$CI_JOB_ID bash -xe $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT


    log_message.status "enroot import: $PYTORCH_IMAGE to /tmp/pytorch.$CI_JOB_ID.enroot.sqsh"
    enroot import --output /tmp/pytorch.$CI_JOB_ID.enroot.sqsh docker://$PYTORCH_IMAGE

    log_message.status "enroot create: pytorch.$CI_JOB_ID"
    enroot create --name pytorch.$CI_JOB_ID /tmp/pytorch.$CI_JOB_ID.enroot.sqsh
    log_message.status "enroot start: pytorch.$CI_JOB_ID"
    enroot start --rw -m /tmp:/tmp pytorch.$CI_JOB_ID bash -xe $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT

    log_message.status "enroot import: $TENSORRT_IMAGE to /tmp/tensorrt.$CI_JOB_ID.enroot.sqsh"
    enroot import --output /tmp/tensorrt.$CI_JOB_ID.enroot.sqsh docker://$TENSORRT_IMAGE
    log_message.status "enroot create: tensorrt.$CI_JOB_ID"
    enroot create --name tensorrt.$CI_JOB_ID /tmp/tensorrt.$CI_JOB_ID.enroot.sqsh
    log_message.status "enroot start: tensorrt.$CI_JOB_ID"
    enroot start --rw -m /tmp:/tmp tensorrt.$CI_JOB_ID bash -xe $TRITON_MDLS_SRC_DIR/$TRTSCRIPT

else
    log_message.warning "Neither Docker nor NVIDIA Enroot is installed." ;
    log_message.warning "Please install Docker or NVIDIA Enroot to generate the models." ;
fi
